// This file was added for video encoding/decoding
#include "QikMovie.h"
#include "RunlengthDecoder.h"
#include "RunlengthEncoder.h"

const char *QikMovie::VERSIONHEADER = "QIM1";

const double QikMovie::defaultCompressionFactor = 1;

QikMovie::QikMovie(): gopsize(1), motionCompensation(true)
{
	SetCompressionFactor(defaultCompressionFactor);
	SetWindowSize(defaultWindowSize);
}

QikMovie::QikMovie(unsigned int newWidth, unsigned int newHeight): gopsize(1), motionCompensation(true)
{
	SetDimension(newWidth, newHeight);
	SetCompressionFactor(defaultCompressionFactor);
	SetWindowSize(defaultWindowSize);
}

QikMovie::~QikMovie()
{
}

void QikMovie::DisableMotionCompensation()
{
	motionCompensation = false;
}

void QikMovie::SetCompressionFactor(double factor)
{
	if (factor < 0)
		factor = 0;

	compressionFactor = factor;
}

void QikMovie::SetWindowSize(unsigned int newWindowsize)
{
	if (windowsize < MACROBLOCKW)
		windowsize = MACROBLOCKW;

	if (windowsize < MACROBLOCKH)
		windowsize = MACROBLOCKH;

	windowsize = newWindowsize;
}

void QikMovie::WriteHeaderToStream(ByteStream& stream)
{
	unsigned int width = Width();
	unsigned int height = Height();
	unsigned int framecount = FrameCount();

	// Write QIM file header
	stream.WriteRaw(VERSIONHEADER, 4);
	stream.WriteRaw(&width, sizeof(width));
	stream.WriteRaw(&height, sizeof(height));
	stream.WriteRaw(&framecount, sizeof(framecount));
	stream.WriteRaw(&gopsize, sizeof(gopsize));

	stream.WriteRaw(&compressionFactor, sizeof(compressionFactor));
}

void QikMovie::WriteIFrameToStream(Qik& frameencoder, Image& frame, ByteStream& stream)
{
	frameencoder.ReadRawFromData( frame.ConstData(), frame.Size() );

	// Write I frame to stream, and update raw image with decoded result
	frameencoder.WriteToStream(stream);

	// Update frame with decoded result
	frame = frameencoder;
}

void QikMovie::WritePFrameToStream(Qik& macroencoder, Image& prevframe, Image& frame, ByteStream& stream)
{
	frame.InitializeSegment(MACROBLOCKW, MACROBLOCKH);

	// For each block
	while(1)
	{
		// Get next image block
		ImageSegment *seg = frame.NextSegment();

		if (!seg)
			break;

		// Find seg in prevframe
		ImageSegment *refseg = seg->LocateSegment(prevframe, windowsize);

		if (!refseg)
			throw "QikMovie: ImageSegment not found in reference frame";

		// Calculate difference
		seg->Substract(*refseg);

		macroencoder.ReadRawFromSegment(0, 0, *seg);

		signed short motionX = refseg->OriginX() - seg->OriginX();
		signed short motionY = refseg->OriginY() - seg->OriginY();

		stream.WriteRaw(&motionX, sizeof(motionX));
		stream.WriteRaw(&motionY, sizeof(motionY));

		// Encode the difference, and write back decoded result
		macroencoder.WriteToStream(stream);

		// Read diff data into segment, which points to frame
		seg->ReadRawFromImage(macroencoder);

		// Reconstruct decoded segment by adding the difference
		seg->Add(*refseg);
	}
}

void QikMovie::ReadIFrameFromStream(Qik& framedecoder, Image& frame, ByteStream& stream)
{
	framedecoder.ReadFromStream(stream);

	frame.SetDimension(Width(), Height());
	frame.ReadRawFromData(framedecoder.ConstData(), framedecoder.Size());
}

void QikMovie::ReadPFrameFromStream(Qik& macrodecoder, Image& prevframe, Image& frame, ByteStream& stream)
{
	frame.SetDimension(Width(), Height());
	frame.InitializeSegment(MACROBLOCKW, MACROBLOCKH);

	// For each block
	while(1)
	{
		// Get next image block
		ImageSegment *seg = frame.NextSegment();

		if (!seg)
			break;

		signed short motionX;
		signed short motionY;

		stream.ReadRaw(&motionX, sizeof(motionX));
		stream.ReadRaw(&motionY, sizeof(motionY));


		// Read diff
		macrodecoder.ReadFromStream(stream);


		seg->ReadRawFromImage(macrodecoder);
		ImageSegment *refseg = prevframe.SegmentFromOffset(seg->OriginX()+motionX, seg->OriginY()+motionY, MACROBLOCKW, MACROBLOCKH);

		if (motionCompensation)
			seg->Add(*refseg);
		else
			seg->ReadRawFromSegment(*refseg);
	}
}


void QikMovie::WriteToFile(const char* filename)
{
	RunLengthEncoder outstream;
	Qik frameencoder(true);
	Qik macroencoder(true);



	macroencoder.SetDimension(MACROBLOCKW, MACROBLOCKH);
	frameencoder.SetDimension(Width(), Height());

	macroencoder.SetCompressionFactor(compressionFactor);
	frameencoder.SetCompressionFactor(compressionFactor);

	// Enable encoded images to update itself with the decoded result
	macroencoder.EnableWriteback();
	frameencoder.EnableWriteback();
	



	WriteHeaderToStream(outstream);

	unsigned int framenr = 0;
	while(!EndOfMovie())
	{
		printf("\rFrame %d / %d", framenr+1, FrameCount());

		if ((framenr%gopsize) == 0)
			WriteIFrameToStream(frameencoder, *CurrentFrame(), outstream);
		else
			WritePFrameToStream(macroencoder, *PreviousFrame(), *CurrentFrame(), outstream);

		SeekNext();
		framenr++;
	}
	printf("\n");

	outstream.WriteToFile(filename);
}


void QikMovie::ReadHeaderFromStream(ByteStream& stream)
{
	char versionString[5];
	unsigned int width, height, framecount;


	memset(versionString, 0, sizeof(versionString));
	stream.ReadRaw(versionString, 4);

	if (strcmp(versionString, VERSIONHEADER))
		throw "QikMovie: Reading invalid qik format";


	stream.ReadRaw(&width, sizeof(width));
	stream.ReadRaw(&height, sizeof(height));
	stream.ReadRaw(&framecount, sizeof(framecount));
	stream.ReadRaw(&gopsize, sizeof(gopsize));
	stream.ReadRaw(&compressionFactor, sizeof(compressionFactor));

	SetDimension(width, height);
	SetFrameCount(framecount);
}

void QikMovie::ReadFromFile(const char* filename)
{
	RunLengthDecoder instream;
	Qik framedecoder(true);
	Qik macrodecoder(true);



	instream.ReadFromFile(filename);

	ReadHeaderFromStream(instream);

	macrodecoder.SetDimension(MACROBLOCKW, MACROBLOCKH);
	framedecoder.SetDimension(Width(), Height());

	macrodecoder.SetCompressionFactor(compressionFactor);
	framedecoder.SetCompressionFactor(compressionFactor);

	unsigned int framenr = 0;
	while(!EndOfMovie())
	{
		printf("\rFrame %d / %d", framenr+1, FrameCount());
		Image *frame = CurrentFrame();

		if ((framenr%gopsize) == 0)
			ReadIFrameFromStream(framedecoder, *CurrentFrame(), instream);
		else
			ReadPFrameFromStream(macrodecoder, *PreviousFrame(), *CurrentFrame(), instream);

		SeekNext();
		framenr++;
	}
	printf("\n");
}

void QikMovie::SetGopSize(unsigned int newGopSize)
{
	if (gopsize)
		gopsize = newGopSize;
}
